import os
import cPickle
import numpy as np
from simple_tif import array_to_tif, tif_to_array

illumination_directory = os.path.join(os.getcwd(), 'illumination')

"""
Separable or not, these parameters determine the illumination:
"""
num_rotations = 3
num_phases = 5
illumination_wavelength = 7.913939 #in estimate pixel units
A1, A2, A3 = 1, 1, 1 #Relative amplitudes of the three beams
beam_tilt_angle = 0.8636 #Radians, angle of the two tilted beams
#Initial phases of the three beams, for each orientation:
phi1 = (6.12177 - 0.5*0.0320431756664,
        -1.7178449999999998 + 0.5*0.368955574059,
        (0.7199 + 1.1598 - 0.98942 - 0.110082367703)*0.5)
phi2 = (-phi1[0], -phi1[1], -phi1[2])
phi3 = (0, 0, 0)
phase_step_size = (1./0.99367, -1./0.999706, -1./0.973466)
##intensity_vs_orientation = (1., 597.596/667.429, 579.7/667.429)
initial_orientation_angle = 0.793944964
#Radians, initial angle of the illumination stripes in the x/y plane.
z_pix_size_over_x_pix_size = 125.0/40.0
#In estimate pixel units; z-pixels are typically bigger than x/y pixels.

def generate_illumination(
    density_shape,
    zoom_factor_xy,
    zoom_factor_z,
    pad_amount_xy=0,
    pad_amount_z=0,
    which_depth='all',
    which_pattern='all',
    check_directory=True
    ):
    """
    For the general case, where the illumination might change every
    time we move the objective focus depth.

    'density_shape' is the shape of the array we'll use to hold our
    estimate of the sample density. The illumination array needs to be
    this shape, too.

    'zoom_factor_xy' is the ratio between the lateral pixel size of
    our estimate of the sample density, and the lateral pixel size of
    the measured data.

    'zoom_factor_z' is the ratio between the axial pixel size of our
    estimate of the sample density, and the axial pixel size of the
    measured data.

    'pad_amount_xy' and 'pad_amount_z' are the number of extra pixels
    we pad the edge of our estimate with in the lateral and axial
    directions, to allow for fluorophores outside the measured volume
    casting light into the measured volume.

    'which_depth'  and 'which_pattern' are currently un-implemented.

    Saves a bunch of data to disk, one 3D tif for every illumination
    pattern, at every depth the objective looks at.
    """
    """
    First off, make sure there's a folder to put the illumination data,
    and make sure the folder is empty:
    """
    if check_directory:
        check_illumination_directory()

    """
    Pre-calculate a few helpful variables
    """
    k = 2 * np.pi / illumination_wavelength
    depth_list = range(pad_amount_z,
                       density_shape[0] - pad_amount_z,
                       zoom_factor_z)
    x = np.arange(density_shape[2]).reshape(1, 1, density_shape[2])
    y = np.arange(density_shape[1]).reshape(1, density_shape[1], 1)
    z = np.arange(density_shape[0]).reshape(density_shape[0], 1, 1)
    """
    Now, make the illumination patterns:
    """
    for d, depth in enumerate(depth_list): #For every depth we focus to
        print "Generating illumination at depth", depth, "",
        print "(", d+1, "of", len(depth_list), ")"
        print " Position:",
        for rot, orientation_angle in enumerate(
            np.arange(initial_orientation_angle,
                      initial_orientation_angle + 2*np.pi,
                      2*np.pi/num_rotations)):
            k1 = k * np.array((
                np.sin(beam_tilt_angle) * np.cos(orientation_angle),
                np.sin(beam_tilt_angle) * np.sin(orientation_angle),
                np.cos(beam_tilt_angle) * z_pix_size_over_x_pix_size))
            k2 = k * np.array((
                -np.sin(beam_tilt_angle) * np.cos(orientation_angle),
                -np.sin(beam_tilt_angle) * np.sin(orientation_angle),
                np.cos(beam_tilt_angle) * z_pix_size_over_x_pix_size))
            k3 = k * np.array((0, 0, z_pix_size_over_x_pix_size))
            for pha, phase in enumerate(
                np.arange(0, 2*np.pi, 2*np.pi/num_phases)):
                p = pha + rot*num_phases
                print p,
                field = (
                    A1 * np.exp(1j * (k1[0] * x +
                                      k1[1] * y +
                                      k1[2] * (z - depth) +
                                      (phi1[rot] +
                                       phase * phase_step_size[rot]))
                                ) + 
                    A2 * np.exp(1j * (k2[0] * x +
                                      k2[1] * y +
                                      k2[2] * (z - depth) +
                                      (phi2[rot] -
                                       phase * phase_step_size[rot]))
                                ) +
                    A3 * np.exp(1j * (k3[0] * x +
                                      k3[1] * y +
                                      k3[2] * (z - depth) +
                                      phi3[rot])))
                illumination = (field * np.conj(field)).real #* intensity_vs_orientation[rot]
                file_name = os.path.join(
                    illumination_directory,
                    'illumination_depth%05i_pattern%05i.tif'%(d, p))
                array_to_tif(illumination.astype(np.float32), file_name)
        print
    set_illumination_info({
        'density_shape': density_shape,
        'depth_list': depth_list,
        'zoom_factor_xy': zoom_factor_xy,
        'zoom_factor_z': zoom_factor_z,
        'num_scan_positions': p + 1,
        'num_depths': len(depth_list),
        'pad_amount_xy': pad_amount_xy,
        'pad_amount_z': pad_amount_z,
        'separable': False,
        'estimate_z_pix_size_over_x_pix_size': z_pix_size_over_x_pix_size,
        })
    return None

def generate_separable_illumination(
    density_shape,
    emission_psf_shape,
    zoom_factor_xy,
    zoom_factor_z,
    pad_amount_xy=0,
    pad_amount_z=0,
    which_depth='all',
    which_pattern='all',
    check_directory=True,
    output_full_illumination=False,
    ):
    """
    For the specific case where the illumination can be written as the
    sum of a finite number of terms, and each term can be written as
    the product of a function only of 'z', and a function only of 'x'
    and 'y'.

    'density_shape' is the shape of the array we'll use to hold our
    estimate of the sample density. The lateral illumination array needs to
    match the x/y dimensions of this shape.

    'emission_psf_shape' is the shape of the array we'll use to hold
    our emission point spread function. The axial illumination array
    needs to match the z dimension of this shape.

    'zoom_factor_xy' is the ratio between the lateral pixel size of
    our estimate of the sample density, and the lateral pixel size of
    the measured data.

    'zoom_factor_z' is the ratio between the axial pixel size of our
    estimate of the sample density, and the axial pixel size of the
    measured data.

    'pad_amount_xy' and 'pad_amount_z' are the number of extra pixels
    we pad the edge of our estimate with in the lateral and axial
    directions, to allow for fluorophores outside the measured volume
    casting light into the measured volume.

    'which_depth'  and 'which_pattern' are currently un-implemented.

    'output_full_illumination' generates separable illumination, but
    saves illumination to disk as non-separable. There's no practical
    reason to do this, but it's nice for debugging, to verify that
    generate_illumination and generate_separable_illumination have
    equivalent output.

    Saves data to disk, a stack of 2D tifs and a stack of 1D tifs, one
    entry in the stack for every illumination pattern.
    """
    """
    First off, make sure there's a folder to put the illumination data,
    and make sure the folder is empty:
    """
    if check_directory:
        check_illumination_directory()

    """
    Pre-calculate a few helpful variables
    """
    k = 2 * np.pi / illumination_wavelength
    kz = k * np.cos(beam_tilt_angle)
    depth_list = range(pad_amount_z,
                       density_shape[0] - pad_amount_z,
                       zoom_factor_z)
    x = np.arange(density_shape[2]).reshape(1, 1, density_shape[2])
    y = np.arange(density_shape[1]).reshape(1, density_shape[1], 1)
    z = np.arange(density_shape[0]).reshape(density_shape[0], 1, 1)
    psf_z = np.arange(emission_psf_shape[0]) - emission_psf_shape[0] // 2
    illumination_z = np.zeros((emission_psf_shape[0], 3))
    illumination_z[:, 0] = 1
    illumination_z[:, 1] = np.cos((k - kz) * psf_z * z_pix_size_over_x_pix_size)
    illumination_z[:, 2] = np.sin((k - kz) * psf_z * z_pix_size_over_x_pix_size)
    illumination_z = illumination_z.reshape(illumination_z.shape + (1,))
    """
    Now, make the illumination patterns:
    """
    for rot, orientation_angle in enumerate(
        np.arange(initial_orientation_angle,
                  initial_orientation_angle + 2*np.pi,
                  2*np.pi/num_rotations)):
        kx = k * np.sin(beam_tilt_angle) * np.cos(orientation_angle)
        ky = k * np.sin(beam_tilt_angle) * np.sin(orientation_angle)
        for pha, phase in enumerate(
            np.arange(0, 2*np.pi, 2*np.pi/num_phases)):
            p = pha + rot*num_phases
            F = (A1**2 + A2**2 + A3**2 +
                 (2*A1*A2 * np.exp(1j * (2 * kx * x +
                                         2 * ky * y +
                                         phi1[rot] - phi2[rot] +
                                         2 * phase * phase_step_size[rot]))
                  ).real)
            G = (A1*A3 * np.exp(+1j * (kx * x +
                                       ky * y +
                                       phi1[rot] - phi3[rot] +
                                       phase * phase_step_size[rot])) +
                 A2*A3 * np.exp(-1j * (kx * x +
                                       ky * y +
                                       -phi2[rot] + phi3[rot] +
                                       phase * phase_step_size[rot])))
            if output_full_illumination: #For debugging
                """
                """
                for d, depth in enumerate(depth_list):
                    illumination = (
                        F +
                        2*G.real * np.cos((k - kz) *
                                          (z - depth) *
                                          z_pix_size_over_x_pix_size) +
                        2*G.imag * np.sin((k - kz) *
                                          (z - depth) *
                                          z_pix_size_over_x_pix_size)
                        ) #* intensity_vs_illumination[rot]
                    file_name = os.path.join(
                        illumination_directory,
                        'illumination_depth%05i_pattern%05i.tif'%(d, p))
                    array_to_tif(illumination.astype(np.float32), file_name)
            else:
                illumination_xy = np.concatenate(
                    (F, 2*G.real, 2*G.imag), axis=0
                    ) #* intensity_vs_orientation[rot]
                file_name = os.path.join(
                    illumination_directory,
                    'illumination_xy_pattern%05i.tif'%(p))
                print "Saving", file_name
                array_to_tif(illumination_xy.astype(np.float32), file_name)
                file_name = os.path.join(
                    illumination_directory,
                    'illumination_z_pattern%05i.tif'%(p))
                print "Saving", file_name
                array_to_tif(illumination_z.astype(np.float32), file_name)
    print "Done generating illumination."
    set_illumination_info({
        'density_shape': density_shape,
        'depth_list': depth_list,
        'zoom_factor_xy': zoom_factor_xy,
        'zoom_factor_z': zoom_factor_z,
        'num_scan_positions': p + 1,
        'num_depths': len(depth_list),
        'pad_amount_xy': pad_amount_xy,
        'pad_amount_z': pad_amount_z,
        'separable': not output_full_illumination,
        'estimate_z_pix_size_over_x_pix_size': z_pix_size_over_x_pix_size,
        })
    return None

def check_illumination_directory():
    if os.path.exists(illumination_directory):
        response = raw_input("Illumination directory exists! Delete? ([Y]/N):")
        if response == 'N':
            return None
        else:
            print "Deleting..."
            files = sorted(os.listdir(illumination_directory))
            for f in files:
                print "Deleting:", f
                os.remove(os.path.join(illumination_directory, f))
            print "Done deleting."
    else:
        os.mkdir(illumination_directory)
    return None

def get_illumination_info():
    illumination_info = cPickle.load(
        open(os.path.join(illumination_directory,
                          'illumination_info.pkl'),
             'rb'))
    return illumination_info

def set_illumination_info(illumination_info):
    cPickle.dump(
        illumination_info,
        open(os.path.join(illumination_directory,
                          'illumination_info.pkl'),
             'wb'))
    return None

def get_illumination_pattern(which_depth, which_pattern):
    file_name = os.path.join(
        illumination_directory,
        'illumination_depth%05i_pattern%05i.tif'%(which_depth, which_pattern))
    return tif_to_array(file_name)

def set_illumination_pattern(new_pattern, which_depth, which_pattern):
    file_name = os.path.join(
        illumination_directory,
        'illumination_depth%05i_pattern%05i.tif'%(which_depth, which_pattern))
    array_to_tif(new_pattern, file_name)
    return None

def get_separable_illumination_pattern(which_pattern, direction='xy'):
    if direction not in ('xy', 'z'):
        print "Direction:", direction
        raise UserWarning("'direction' must be 'xy' or 'z'")
    file_name = os.path.join(
        illumination_directory,
        'illumination_' + direction + '_pattern%05i.tif'%(which_pattern))
    return tif_to_array(file_name)

def set_separable_illumination_pattern(
    new_pattern, which_pattern, direction='xy'):
    if direction not in ('xy', 'z'):
        print "Direction:", direction
        raise UserWarning("'direction' must be 'xy' or 'z'")
    file_name = os.path.join(
        illumination_directory,
        'illumination_' + direction + '_pattern%05i.tif'%(which_pattern))
    array_to_tif(new_pattern, file_name)
    return None

if __name__ == '__main__':
    illumination_directory=os.path.join(os.getcwd(), 'illumination_old')
    generate_illumination(
        density_shape=(40, 60, 60),
        zoom_factor_xy=2,
        zoom_factor_z=2)
    print get_illumination_info()

    illumination_directory=os.path.join(os.getcwd(), 'illumination_new')
    generate_separable_illumination(
        density_shape=(40, 60, 60),
        emission_psf_shape=(40, 60, 60),
        zoom_factor_xy=2,
        zoom_factor_z=2,
        output_full_illumination=True)
    print get_illumination_info()
    
